#include <cstdio>
#include <cstdlib>
#include <unistd.h>
#include <iostream>

#include "wave_reader.h"
#include "../src/webrtc_aec3.h"

constexpr char* DEFAULT_PLAY_AUDIO_FILE = "lay.lady.lay.wav";
constexpr char* DEFAULT_RECORD_AUDIO_FILE = "mix_audio.wav";
constexpr char* DEFAULT_OUTPUT_AUDIO_FILE = "separated_speech.wav";
constexpr int DEFAULT_SAMPLE_RATE = 48000;
constexpr int DEFAULT_CHANNEL_COUNT = 2;
constexpr int DEFAULT_SAMPLE_LENGTH_MS = 50;

struct InputConfigure
{
    char* input_play_auio_file_name_;
    char* input_record_audio_file_name_;
    char* output_file_name_;
};

void DumpInputInfo(const InputConfigure& cfg)
{
    std::cout   << "input info:\n" 
                << "    input record file: " << cfg.input_record_audio_file_name_ << "\n"
                << "    input play file: " << cfg.input_play_auio_file_name_ << "\n"
                << "    specified output file: " << cfg.output_file_name_ << "\n"
                << std::endl;
}

void Usage(const char* program_name) 
{
    std::cout << program_name << ":\n" 
              << "          -i input record audio file name. default: " << DEFAULT_RECORD_AUDIO_FILE << "\n"
              << "          -p input play audio file name. default: " << DEFAULT_PLAY_AUDIO_FILE << "\n" 
              << "          -o output file name. default: " << DEFAULT_OUTPUT_AUDIO_FILE <<  "\n" 
              << "          -r sample rate, default: " << DEFAULT_SAMPLE_RATE << "\n" 
              << "          -c channel count, default: " << DEFAULT_CHANNEL_COUNT << "\n" 
              << "          all audio should save in wav format, encode as pcm16le\n"
              <<std::endl;
}

void PrintError(char* info)
{
    fprintf(stderr, "[%s][%d][%s]: %s\n", __FILE__, __LINE__, __FUNCTION__, info);
}

int main(int argc, char* argv[])
{
    int opt = -1;
    struct InputConfigure cfg = { DEFAULT_PLAY_AUDIO_FILE, 
                                    DEFAULT_RECORD_AUDIO_FILE, 
                                    DEFAULT_OUTPUT_AUDIO_FILE };

    // get option
    while((opt = getopt(argc, argv, "i::p::o::r::c::h")) != -1) {
        switch(opt) {
        case 'i':
            cfg.input_record_audio_file_name_ = optarg;
        case 'p':
            cfg.input_play_auio_file_name_ = optarg;
        case 'o':
            cfg.output_file_name_ = optarg;
        default:
            Usage(argv[0]);
            break;
        }
    }

    // dump_input_info
    DumpInputInfo(cfg);

    // open file
    WaveReader record_file_reader;
    WaveReader play_file_reader;
    WaveWriter output_file_writer;

    if(record_file_reader.Open(cfg.input_record_audio_file_name_) != 0) {
        PrintError("open record file failed");
        exit(-1);
    }

    if(play_file_reader.Open(cfg.input_play_auio_file_name_) != 0) {
        PrintError("open play file failed");
        exit(-1);
    }

    if(output_file_writer.Open(cfg.output_file_name_) != 0) {
        PrintError("open output file failed");
        exit(-1);
    }

    // get record/play file audio info, they should be the sample
    uint32_t record_audio_sample_rate = 0;
    uint16_t record_audio_channel_count = 0;
    uint16_t record_audio_bit_per_sample = 0;
    record_file_reader.GetAudioInfo(record_audio_sample_rate, record_audio_channel_count, record_audio_bit_per_sample);

    uint32_t play_audio_sample_rate = 0;
    uint16_t play_audio_channel_count = 0;
    uint16_t play_audio_bit_per_sample = 0;
    play_file_reader.GetAudioInfo(play_audio_sample_rate, play_audio_channel_count, play_audio_bit_per_sample);

    if(record_audio_sample_rate != play_audio_sample_rate ||
    record_audio_channel_count != play_audio_channel_count ||
    record_audio_bit_per_sample != play_audio_bit_per_sample) {
        PrintError("play and record audio property should be sample");
        exit(-1);
    }

    // setup output file audio info
    output_file_writer.SetAudioInfo( record_audio_sample_rate, record_audio_channel_count, record_audio_bit_per_sample);

    // init aec3 struct
    int sample_per_frame = play_audio_sample_rate *  DEFAULT_SAMPLE_LENGTH_MS / 1000;     
    int sample_size = play_audio_sample_rate * play_audio_channel_count * (play_audio_bit_per_sample/8) * DEFAULT_SAMPLE_LENGTH_MS / 1000;
    int echo_option = WEBRTC_AEC_OPTION_USE_NOISE_SUPPRESSOR | WEBRTC_AEC_OPTION_USE_GAIN_CONTROLLER;
    void* aec3_config = NULL;

    int ret = webrtc_aec3_create(&aec3_config, play_audio_sample_rate, play_audio_channel_count, sample_per_frame, 0, echo_option);
    if(ret != WEBRTC_AEC_ERROR_SUCCESS) {
        PrintError("webrtc_aec3_create failed");
        exit(0);
    }

    // read audio and apply the audio echo cancel
    std::shared_ptr<char[]> play_audio_buf(new char[sample_size]);
    std::shared_ptr<char[]> record_audio_buf(new char[sample_size]);

    std::size_t pnread, rnread;
    // apply aec to audio frame
    while(true) {
        pnread = play_file_reader.Read(play_audio_buf.get(), sample_size);
        if(pnread <= 0) {
            break;
        }
        rnread = record_file_reader.Read(record_audio_buf.get(), sample_size);

        if(WEBRTC_AEC_ERROR_SUCCESS != webrtc_aec3_cancel_echo(aec3_config, (int16_t*)record_audio_buf.get(), (const int16_t*)play_audio_buf.get())) {
            PrintError("cancel echo failed");
        }

        // write to file
        if(1 != output_file_writer.Write(record_audio_buf.get(), sample_size, 1)) {
            PrintError("write audio failed");
        }
    }

    // destroy aec3 struct
    webrtc_aec3_destroy(aec3_config);
    
    output_file_writer.Close();

    exit(0);
}